package org.honeynet.hbbackend.shadowserverdns;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.Arrays;
import java.util.EnumSet;
import java.util.Hashtable;
import java.util.regex.Pattern;

import javax.annotation.Resource;
import javax.ejb.EJBException;
import javax.ejb.Schedule;
import javax.ejb.SessionContext;
import javax.ejb.Singleton;
import javax.ejb.TransactionAttribute;
import javax.ejb.TransactionAttributeType;
import javax.jms.ConnectionFactory;
import javax.jms.JMSException;
import javax.jms.Message;
import javax.jms.Session;
import javax.jms.Topic;
import javax.jms.TopicSubscriber;
import javax.naming.Context;
import javax.naming.NamingException;
import javax.naming.directory.Attributes;
import javax.naming.directory.DirContext;
import javax.naming.directory.InitialDirContext;
import javax.sql.DataSource;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@Singleton
public class ShadowServerDnsBinding {

	private final static String SERVER = "asn.shadowserver.org";
	private final static Pattern SPLIT_REGEX = Pattern.compile("\\s+\\|\\s+");
	private final static int RECV_TIMEOUT = 3000; 
	
	private static Logger log = LoggerFactory.getLogger(ShadowServerDnsBinding.class);
	
	@Resource(mappedName="jdbc/hbbackend")
	DataSource ds;
	
	@Resource(mappedName="jms/DurableConsumer/shadowserverdns")
	private ConnectionFactory connectionFactory;
	
	@Resource(mappedName="jms/new_ip")
	private Topic dest;

	@Resource
	private SessionContext sctx;
	
	private boolean stop = false;

	
	@Schedule(second="*/5", minute="*", hour="*", persistent=false)
	@TransactionAttribute(TransactionAttributeType.REQUIRES_NEW)
	public void work() {
		log.trace("work() called");
		
		javax.jms.Connection mq = null;
		Session sess = null;
		TopicSubscriber sub = null;
		
		try {
			mq = connectionFactory.createConnection();
			sess = mq.createSession(true, 0);
			sub = sess.createDurableSubscriber(dest, "shadowserverdns");
			log.trace("created durable subscription");
			log.trace("starting receive loop");
			mq.start();
			
			Message jmsMsg;
			do {
				if (stop) {
					log.trace("encountered stop request: returning");
					return;
				}
				log.trace("receive");
				jmsMsg = sub.receive(RECV_TIMEOUT);
			} while (jmsMsg == null);
			
			//TODO message validation as the service is ipv4 only(?)

			String ip = jmsMsg.getStringProperty("ip");
			log.debug("received new message ip = {}", ip);

			store(ip, retrieve(ip));
			return;
			
		} catch (JMSException e) {
			log.error("got JMSException", e);
			throw new EJBException(e);
		} finally {
			try {
//				if (sub != null) sub.close();
//				if (sess != null) sess.close();
				if (mq != null) mq.close();
			} catch (JMSException e) {
				log.error("got JMSException while trying to close", e);
				throw new EJBException(e);
			}
		}
		
	}
	
	
	private Result retrieve(String ip) {
		
		String emsg;
		
		DirContext ctx = null;
		
		
		String[] octets = ip.split("\\.");
		StringBuilder sb = new StringBuilder(22);
		for (int i = octets.length -1 ; i >= 0 ; i--) {
			sb.append(octets[i]);
			sb.append(".");
		}
		sb.append("origin");

		String req = sb.toString();
		
		try {
			
			Hashtable<String, String> env =	new Hashtable<String, String>();
			env.put(Context.INITIAL_CONTEXT_FACTORY, "com.sun.jndi.dns.DnsContextFactory");	
			env.put(Context.PROVIDER_URL, "dns://" + SERVER + "/shadowserverdns");
			
			log.trace("querying {} req = {}", SERVER, req);
			ctx = new InitialDirContext(env);
			Attributes attrs = ctx.getAttributes(req, new String[] {"TXT"});
			
			String response = attrs.get("TXT").toString();
			// strip "TXT: "
			try {
				response = response.substring("TXT: ".length());
			} catch (IndexOutOfBoundsException e) {
				emsg = "couldn't parse return value: " + response;
				log.error(emsg);
				throw new EJBException(emsg);
			}
			
			String[] responseFields = SPLIT_REGEX.split(response);
			log.trace("parsed response fields = {}", Arrays.asList(responseFields));
			
			// general check
			if (responseFields.length != Fields.values().length) {
				emsg = "invalid number of fields in response";
				log.error(emsg);
				throw new EJBException(emsg);
			}
			
			// general validation
			for (Fields field : EnumSet.allOf(Fields.class)) {
				String value = responseFields[field.ordinal()].trim();
				log.trace("read field: {} = {}", field.name(), value);
				if (field.validate(value) == false) {
					emsg = "field does not validate: " + field.name();
					log.error(emsg);
					throw new EJBException(emsg);
				}
			}
			
			Result res = new Result();
			res.asn = Long.parseLong(responseFields[Fields.asn.ordinal()]);
			res.as_name = responseFields[Fields.as_name.ordinal()];
			res.cc = responseFields[Fields.cc.ordinal()];
			res.dom = responseFields[Fields.dom.ordinal()];
			res.isp = responseFields[Fields.isp.ordinal()];
			res.bgp_prefix = responseFields[Fields.bgp_prefix.ordinal()];
			
			log.trace("result: {}", res);
			
			return res;

		} catch (NamingException e) {
			log.error("got NamingException", e);
			throw new EJBException(e);
		} finally {
			try {
				if (ctx != null) ctx.close();
			} catch (NamingException e) {
				log.error("got NamingException", e);
				throw new EJBException(e);
			}
		}
	}
	
	
	private void store(String ip, Result res) {
		log.trace("store() called");

		Connection db = null;
		PreparedStatement pStmt;

		try {
			db = ds.getConnection();
			pStmt = db.prepareStatement("insert into asn_shadowserver (ip, asn, as_name, cc, dom, isp, bgp_prefix) values (inet(?),?,?,?,?,?,inet(?))");

			log.trace("inserting into asn_shadowserver (ip = {})", ip);

			pStmt.setString(1, ip);
			pStmt.setLong(2, res.asn);
			pStmt.setString(3, res.as_name);
			pStmt.setString(4, res.cc);
			pStmt.setString(5, res.dom);
			pStmt.setString(6, res.isp);
			pStmt.setString(7, res.bgp_prefix);

			pStmt.executeUpdate();
			pStmt.close();

		} catch (SQLException e) {
			log.error("got SQLException", e);
			throw new EJBException(e);
		} finally {
			try {
				if (db != null) db.close();
			} catch (SQLException e) {
				log.error("got SQLException", e);
				throw new EJBException(e);
			}
		}
	}
	
	
	
	
	private static enum Fields {
		asn,
		bgp_prefix,
		as_name,
		cc,
		dom,
		isp;
		
		public boolean validate(String value) {
			return true;
		}
	}

	
	private static class Result {
		long asn;
		String as_name;
		String cc;
		String dom;
		String isp;
		String bgp_prefix;
		
		public String toString() {
			return
				Fields.asn.name() + " = " + Long.toString(asn) + " | " +
				Fields.as_name.name() + " = " + as_name + " | " +
				Fields.cc.name() + " = " + cc + " | " +
				Fields.dom.name() + " = " + dom + " | " +
				Fields.isp.name() + " = " + isp+ " | " +
				Fields.bgp_prefix.name() + " = " + bgp_prefix;
		}
	}

}
